package com.yuxl.ai.chat.bailian.service;

import com.yuxl.ai.chat.bailian.domain.BailianReRankInput;
import com.yuxl.ai.chat.bailian.domain.BailianReRankParameters;
import com.yuxl.ai.chat.bailian.domain.BailianReRankReq;
import com.yuxl.ai.chat.bailian.domain.BailianReRankRes;
import com.yuxl.ai.chat.config.ReRanProperties;
import com.yuxl.ai.chat.domain.ReRankReq;
import com.yuxl.ai.chat.domain.ReRankRes;
import com.yuxl.ai.chat.service.ReRankApi;
import com.yuxl.common.core.utils.HttpUtil;
import com.yuxl.common.core.utils.ObjectUtil;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.http.client.reactive.ReactorClientHttpConnector;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;
import reactor.netty.http.client.HttpClient;

import java.util.List;

/**
 * @Author : yuxl
 * @create 2025/7/11 下午2:46
 */
@Slf4j
@Service
@ConditionalOnProperty(prefix = "ai.sort", name = "type", havingValue = "bailian")
public class BailianReRankApi implements ReRankApi {

    @Resource
    private ReRanProperties reRanProperties;
    @Resource
    private ObjectUtil objectUtil;

    @Override
    public Mono<List<ReRankRes>> sort(ReRankReq req) {
        BailianReRankReq r = new BailianReRankReq()
                .setModel(reRanProperties.getModel())
                .setInput(
                        new BailianReRankInput()
                                .setQuery(req.getQuery())
                                .setDocuments(req.getTexts())
                )
                .setParameters(
                        new BailianReRankParameters()
                                .setReturnDocuments(true)
                                .setTopN(req.getTopK())
                );
        return api(r).map(res ->
                res.getOutput().getResults().stream()
                        .map(o -> new ReRankRes()
                                .setText(o.getDocument().getText())
                                .setScore(o.getRelevanceScore()))
                        .toList()
        );
    }

    public Mono<BailianReRankRes> api(BailianReRankReq body) {
        HttpClient httpClient = HttpUtil.httpClientSetTimeout(1000, Long.MAX_VALUE, Long.MAX_VALUE);
        String b = objectUtil.writeValueAsString(body);
        log.debug("sort-api=>body:{}",b);
        return WebClient
                .builder()
                .codecs(configurer ->
                        configurer.defaultCodecs().maxInMemorySize(20 * 1024 * 1024) // 设置为 20MB
                )
                .clientConnector(new ReactorClientHttpConnector(httpClient))
                .baseUrl(reRanProperties.getBaseUrl())
                .build()
                .post()
                .header("content-type", "application/json")
                .header("Authorization", "Bearer " + reRanProperties.getApiKey())
                .bodyValue(body)
                .retrieve()
                .bodyToMono(String.class)
                .map(r -> {
                    log.debug("sort-api=>res:{}",r);
                    return objectUtil.readValue(r, BailianReRankRes.class);
                })
                ;
    }
}
